from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from core.connectors.db_connector import DatabaseConnector


class MySQLConnector(DatabaseConnector):
    def __init__(self, host="localhost", user="", password="", database="", port=3306, pool_size=5):

        v = [str(item) for item in [host, user, database, port]]
        un = "_".join(v)
        super().__init__(un)
        
        """
        初始化 MySQL 数据库连接池
        :param host: 主机地址
        :param user: 用户名
        :param password: 密码
        :param database: 数据库名
        :param port: 端口号
        :param pool_size: 连接池大小
        """
        self.engine = create_engine(
            f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}",
            pool_size=pool_size,
            max_overflow=10,
            pool_recycle=1800,
            pool_pre_ping=True,
            echo=False
        )
        self.Session = sessionmaker(bind=self.engine)

    def get_db_type(self):
        return 'MySQL'

    def test_connection(self):
        """
        测试数据库连接是否成功。
        :return: True 如果连接成功；False 和错误信息如果连接失败。
        """
        try:
            # 尝试连接并执行简单查询
            with self.engine.connect() as connection:
                connection.execute(text("SELECT 1"))
            return True,"Successful"
        except Exception as e:
            return False, str(e)

    def get_session(self):
        """
        获取 SQLAlchemy 会话
        :return: Session 对象
        """
        return self.Session()

    def fetch_table_schema(self, table):
        """
        使用 SQLAlchemy 获取表的字段信息。
        :param table: 表名
        :return: 表字段信息的列表
        """
        query = text("""
        SELECT COLUMN_NAME, COLUMN_COMMENT, DATA_TYPE
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_NAME = :table
        """)

        with self.get_session() as session:
            result = session.execute(query, {"table": table})
            return [row._asdict() for row in result]

    def get_tables_info(self):
        """
        获取 MySQL 数据库中的表名、记录数和字段个数
        :return: 包含表名、记录数和字段个数的列表
        """
        tables_info = []
        with self.engine.connect() as conn:
            tables_query = text("SHOW TABLES;")
            tables = conn.execute(tables_query).fetchall()
            for table_row in tables:
                table = table_row[0]
                # 获取记录数
                row_count_query = text(f"SELECT COUNT(*) FROM `{table}`;")
                row_count = conn.execute(row_count_query).scalar()
                # 获取字段个数
                column_count_query = text(f"DESCRIBE `{table}`;")
                column_count = len(conn.execute(column_count_query).fetchall())
                tables_info.append((table, row_count, column_count))
        return tables_info

    def get_table_ddl(self, table_name: str) -> str:
        """
        获取指定表的建表脚本 (DDL)
        :param table_name: 表名
        :return: 建表脚本字符串
        """
        with self.engine.connect() as conn:
            ddl_query = text(f"SHOW CREATE TABLE `{table_name}`;")
            result = conn.execute(ddl_query).fetchone()
            if result and len(result) > 1:
                return result[1]  # 第二列为建表脚本
            else:
                raise ValueError(f"无法获取表 {table_name} 的建表脚本")


    def get_table_fields(self, table_name):
        """
        获取指定表的字段信息
        :param table_name: 表名称
        :return: 字段信息列表
        """
        query = text(f"""
        SELECT
            COLUMN_NAME AS FieldName,
            COLUMN_COMMENT AS Comment,
            DATA_TYPE AS DataType
        FROM
            INFORMATION_SCHEMA.COLUMNS
        WHERE
            TABLE_SCHEMA = '{self.engine.url.database}'
            AND TABLE_NAME = :table_name;
        """)

        with self.engine.connect() as connection:
            result = connection.execute(query, {"table_name": table_name}).fetchall()
            return [{"FieldName": row[0], "Comment": row[1], "DataType": row[2], "IS_GEN": False} for row in result]


    def close(self):
        """
        关闭 SQLAlchemy 引擎
        """
        self.engine.dispose()
